/*
 * Functions to solve data flow equations and
 * do the related optimizations.
 */
#include "HIRCode.h"
#include "HXD1.h"
#include <algorithm>

// Don't use these flags of temp!
#define fassigned ERROR
#define frefed ERROR
#define lrefed ERROR

/*
 * Data flow analysis: solve equations of live-variable analysis.
 */
void HIRFunct::dfa_solve_livevar(void)
{
	/*
	 * Calculate "use"(`livevar_use`) and "def"(`livevar_def`)
	 * set for every BB and then solve data flow equations to get
	 * "IN"(`livevar_in`) and "OUT"(`livevar_def`) set for them.
	 * Hint: IN[B] = use[B] + (OUT[B] - def[B]).
	 */
	uint32_t iidx;
	std::vector<HIRTempno> inst_srctemps;

	for(auto bb = basic_blocks.begin(); bb != basic_blocks.end(); ++bb)
	{
		bb->livevar_use.clear();
		bb->livevar_def.clear();
		bb->livevar_in.clear();
		bb->livevar_out.clear();
		// Calculate "use"(`livevar_use`)
		if (bb->begin == INVALID_INST_IDX)
			// skip the #0 BB if it is empty
			continue;
		for (iidx = bb->begin; iidx <= bb->end; ++iidx)
		{
			ircode.at(iidx).get_srctemps(inst_srctemps);
			for (auto st : inst_srctemps)
				if (bb->asstemps.find(st) == bb->asstemps.end()
					|| bb->asstemps[st] >= iidx)
				/*
				 * If the source temp is not assigned in this BB, or is
				 * assigned but after this instruction, then add it into "use".
				 */
					bb->livevar_use.insert(st);
		}

		// Calculate "def"(`livevar_def`)
		for (auto & asstemp : bb->asstemps)
			// assigned temps that are not in `livevar_use`
			if (bb->livevar_use.find(asstemp.first) == bb->livevar_use.end())
				bb->livevar_def.insert(asstemp.first);
	}

	bool in_chgd = true;
	// Calculate "IN"(`livevar_in`) and "OUT"(`livevar_def`) sets
	while (in_chgd)
	{
		in_chgd = false;
		for (auto bb = basic_blocks.begin(); bb != basic_blocks.end(); ++bb)
		{
			// OUT[B] = Union{IN[S] | S is a succ of B}
			bb->livevar_out.clear();
			for (auto s : bb->succ)
				for (auto iv : basic_blocks.at(s).livevar_in)
					bb->livevar_out.insert(iv);

			// IN[B] = use[B] + (OUT[B] - def[B])
			std::set<HIRTempno> new_in, temp_set;
			std::set_difference(bb->livevar_out.begin(), bb->livevar_out.end(),
				bb->livevar_def.begin(), bb->livevar_def.end(),
				std::inserter(temp_set, temp_set.begin())
			);
			std::set_union(bb->livevar_use.begin(), bb->livevar_use.end(),
				temp_set.begin(), temp_set.end(),					 
				std::inserter(new_in, new_in.begin())
			);
			if (bb->livevar_in != new_in)
			{
				in_chgd = true;
				bb->livevar_in = new_in;
			}
		}
	}
}

/*
 * For every reachable basic blocks, calculate their
 * dominators and push their indexes into `dominators`.
 */
void HIRFunct::dfa_solve_dominators(void)
{
	// use iterative algorithm to calculate dominators for every reachable BB.

	// init `BB->dominators`
	std::set<uint32_t> reachable_bbs;
	for (size_t bbidx = 0ul; bbidx < basic_blocks.size(); ++bbidx)
		if (basic_blocks.at(bbidx).reachable)
			reachable_bbs.insert(bbidx);
	for (uint32_t rbbidx : reachable_bbs)
		basic_blocks.at(rbbidx).dominators = reachable_bbs;

	basic_blocks.at(0ul).dominators = {0u};

	// calculate dominators
	bool flag_changed = true;
	while (flag_changed)
	{
		flag_changed = false;
		auto ebb = reachable_bbs.begin();
		for (++ebb; ebb != reachable_bbs.end(); ++ebb)
		{
			// foreach non-zero reachable BB
			std::set<uint32_t> dom = reachable_bbs;
			for (uint32_t pred : basic_blocks.at(*ebb).pred)
			{
				std::set<uint32_t> ndom;
				std::set_intersection(dom.begin(), dom.end(), 
					basic_blocks.at(pred).dominators.begin(),
					basic_blocks.at(pred).dominators.end(),
					std::inserter(ndom, ndom.begin())
				);
				dom = ndom;
			}
			dom.insert(*ebb);
			if (basic_blocks.at(*ebb).dominators != dom)
			{
				basic_blocks.at(*ebb).dominators = dom;
				flag_changed = true;
			}
		}
	}
}

/*
 * After dominators are calculated, detect all loops of the function
 * and push them into `loops`.
 */
void HIRFunct::dfa_detect_loops(void)
{
	struct Backedge {
		uint32_t source, sink;
	};
	std::vector<Backedge> backedges;
	std::vector<uint32_t> stack;
	std::set<uint32_t> loop;
	std::vector<bool> visited(basic_blocks.size(), false);

	/*
	 * Use DFS to traverse the control flow graph,
	 * find all back edges and push them to `backedges`.
	 */
	stack.push_back(0u);
	visited.at(0u) = true;
	while (!stack.empty())
	{
		uint32_t tbb = stack.at(stack.size() - 1ul);
		stack.pop_back();
		for (auto succ : basic_blocks.at(tbb).succ)
		{
			if (visited.at(succ) == true &&
				basic_blocks.at(tbb).dominators.count(succ) > 0ul)
			{	// One back edge is found
				Backedge bkedge = {tbb, succ};
				backedges.push_back(bkedge);
			}
			else if (visited.at(succ) == false)
			{
				stack.push_back(succ);
				visited.at(succ) = true;
			}
		}
	}

	/*
	 * For every back edge (`s`, `t`), do inverse DFS from `t` to `s`
	 * to find the loop, and then push the indexes into `loop`.
	 * Then push `loop` into `loops` of the function.
	 */
	for (auto & bkedge : backedges)
	{
		stack.clear();
		loop.clear();
		loop.insert(bkedge.source);
		loop.insert(bkedge.sink);
		stack.push_back(bkedge.source);
		while (!stack.empty())
		{
			uint32_t pbb = stack.at(stack.size() - 1);
			stack.pop_back();
			for (auto pred : basic_blocks.at(pbb).pred)
			{
				if (basic_blocks.at(pred).reachable && loop.count(pred) == 0ul)
				{
					loop.insert(pred);
					stack.push_back(pred);
				}
			}
		}
		loops.push_back(loop);
	}
}

/*
 * Calculate the spill cost of every temporary
 * (labels and `ALLOCA` temporaries are not included).
 * NOTE: this cost is not relevant to the degree of the
 * 	vertex in the interference graph.
 */
void HIRFunct::dfa_calc_spillcost(void)
{
	std::set<HIRTempno> alloca_dests;
	double weight;
	uint32_t bbidx, begin, end, iidx;
	std::vector<HIRTempno> srctemps;
	HIRTempno dest;

	for (uint32_t i = 0u; i < nalloca; ++i)
		alloca_dests.insert(ircode.at(i).get_dest());
	for (bbidx = 1u; bbidx < basic_blocks.size(); ++bbidx)
	{
		weight = 1.0;
		// The weight is 10^n, where n is the depth of loops
		for (auto loop : loops)
			if (loop.count(bbidx) > 0ul)
				weight *= 10.0;
		begin = basic_blocks.at(bbidx).begin;
		end = basic_blocks.at(bbidx).end;
		for (iidx = begin; iidx <= end; ++iidx)
		{
			ircode.at(iidx).get_srctemps(srctemps);
			for (auto st : srctemps)
				temps.at(st).spillcost += weight;
			dest = ircode.at(iidx).get_dest();
			if (dest != INVALID_TEMP_NO)
				temps.at(dest).spillcost += weight;
		}
	}
}